#include <dialects/Builtin/Types.h>
#include <dialects/Math/Ops.h>
#include <dialects/Memory/Ops.h>
#include <dialects/PL/Ops.h>
#include <dialects/PL/Translation/ToMath.h>
#include <lcir/Dag2Dag/D2DDriver.h>
#include <lcir/Dag2Dag/Pattern.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/IRBuilder.h>
#include <lcir/IR/IRRewritter.h>
#include <lcir/Utils/Result.h>

namespace lcir::pl {

static bool checkI32(Value *v) { return v->getType()->isa<Int32Type>(); }
static bool checkF(Value *v) { return v->getType()->isa<FloatType>(); }
static bool typeTrans(Expression *expr, IRRewritter &b) {
  auto op0 = expr->getOperandValue(0);
  auto op1 = expr->getOperandValue(1);

  if (checkI32(op0) && checkI32(op1)) return true;
  if (checkF(op0) && checkF(op1)) return false;
  if (checkF(op0) && checkI32(op1)) {
    auto f = b.createAndGetRet<math::IToF>({op1});
    expr->getOperand(1)->setValue(f);
  }
  if (checkI32(op0) && checkF(op1)) {
    auto f = b.createAndGetRet<math::IToF>({op0});
    expr->getOperand(0)->setValue(f);
  }
  return true;
}
static bool isOperandIsVar(Expression *expr) {
  for (auto op : expr->getOperands()) {
    auto defExpr = op->getValue()->getDefExpr();
    if (defExpr && defExpr->isOp<pl::DeclVar>()) return true;
  }
  return false;
}
template <class TargetExpr, class TargetFExpr>
static Result exprTrans(Expression *expr, IRRewritter &b) {
  if (isOperandIsVar(expr)) return failure();
  typeTrans(expr, b);
  auto op0 = expr->getOperandValue(0);
  auto op1 = expr->getOperandValue(1);
  Value *newVal{nullptr};
  if (checkI32(op0)) {
    newVal = b.createAndGetRet<TargetExpr>({op0, op1});
  } else {
    newVal = b.createAndGetRet<TargetFExpr>({op0, op1});
  }
  b.replaceValue(expr->getResult(), newVal);
  return success();
}

struct InzPattern : public OpPattern<pl::Inz> {
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    auto val = expr->getOperandValue(0);
    Value *newVal{nullptr};
    if (checkI32(val))
      newVal = b.createAndGetRet<math::IInz>({val});
    else if (checkF(val))
      newVal = b.createAndGetRet<math::FInz>({val});
    if (!newVal) {
      std::cout << "Inz转换失败"
                << "\n";
      return failure();
    }
    b.replaceValue(val, newVal);
    return success();
  }
};

struct NegPattern : public OpPattern<pl::Neg> {
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    auto val = expr->getOperandValue(0);
    Value *newVal{nullptr};
    if (val->getType()->isa<Int32Type>())
      newVal = b.createAndGetRet<math::INeg>({val});
    else if (val->getType()->isa<FloatType>())
      newVal = b.createAndGetRet<math::FNeg>({val});
    if (!newVal) {
      std::cout << "NEG转换失败"
                << "\n";
      return failure();
    }
    b.replaceValue(val, newVal);
    return success();
  }
};
struct SeqzPattern : public OpPattern<pl::Seqz> {
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    auto val = expr->getOperandValue(0);
    Value *newVal{nullptr};
    if (val->getType()->isa<Int32Type>())
      newVal = b.createAndGetRet<math::ISeqz>({val});
    else if (val->getType()->isa<FloatType>())
      newVal = b.createAndGetRet<math::FSeqz>({val});
    if (!newVal) {
      std::cout << "EQZ转换失败"
                << "\n";
      return failure();
    }
    b.replaceValue(val, newVal);
    return success();
  }
};

struct AddPattern : public Pattern {
  AddPattern() { setRootNode(PatternNode(pl::Add::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::Add, math::FAdd>(expr, b); }
};

struct SubPattern : public Pattern {
 public:
  SubPattern() { setRootNode(PatternNode(pl::Sub::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::Sub, math::FSub>(expr, b); }
};
struct MulPattern : public Pattern {
 public:
  MulPattern() { setRootNode(PatternNode(pl::Mul::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::Mul, math::FMul>(expr, b); }
};
struct DivPattern : public Pattern {
 public:
  DivPattern() { setRootNode(PatternNode(pl::Div::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::Div, math::FDiv>(expr, b); }
};
struct ModPattern : public Pattern {
 public:
  ModPattern() { setRootNode(PatternNode(pl::Mod::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    if (isOperandIsVar(expr)) return failure();
    auto op0 = expr->getOperandValue(0);
    auto op1 = expr->getOperandValue(1);
    auto newVar = b.createAndGetRet<math::Mod>({op0, op1});
    b.replaceValue(expr->getResult(), newVar);
    return success();
  }
};
struct GTPattern : public Pattern {
 public:
  GTPattern() { setRootNode(PatternNode(pl::GT::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::GT, math::FGT>(expr, b); }
};
struct GEPattern : public Pattern {
 public:
  GEPattern() { setRootNode(PatternNode(pl::GE::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::GE, math::FGE>(expr, b); }
};
struct LTPattern : public Pattern {
 public:
  LTPattern() { setRootNode(PatternNode(pl::LT::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::LT, math::FLT>(expr, b); }
};
struct LEPattern : public Pattern {
 public:
  LEPattern() { setRootNode(PatternNode(pl::LE::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::LE, math::FLE>(expr, b); }
};
struct EQPattern : public Pattern {
 public:
  EQPattern() { setRootNode(PatternNode(pl::EQ::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::EQ, math::FLE>(expr, b); }
};
struct UEPattern : public Pattern {
 public:
  UEPattern() { setRootNode(PatternNode(pl::UE::getOpInstance())); }
  Result rewrite(Expression *expr, IRRewritter &b) const override { return exprTrans<math::UE, math::FUE>(expr, b); }
};

// 开始加patttern
void ToMath::run() {
  auto *module = getLambda();
  D2DDriver driver(module);
  driver.addPatterns<NegPattern, SeqzPattern, AddPattern, SubPattern, MulPattern, DivPattern, ModPattern>();
  driver.addPatterns<GTPattern, GEPattern, LTPattern, LEPattern, EQPattern, UEPattern>();
  driver.run();
}
}  // namespace lcir::pl